'''
Author: duliang thinktanker@163.com
Date: 2025-09-04 08:19:06
LastEditors: duliang thinktanker@163.com
LastEditTime: 2025-10-08 10:52:25
FilePath: 
Description: 
'''
import time
import datetime
import sys
import os
# import flask
import threading
# 添加myyolo模块所在的目录到Python路径中
# 假设myyolo.py文件在当前脚本的同级目录下
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# 如果myyolo.py在其他位置，请修改下面这行的路径
# sys.path.append(r'path/to/myyolo_directory')

from myyolo import MyYOLO
from flask import Flask, request, jsonify, send_file
import cv2
import numpy as np
# import io
# from PIL import Image

class YOLODetector:

    def __init__(self, model_path, image_path, output_folder):
        self.model_path = model_path
        self.image_path = image_path
        self.output_folder = output_folder
        self.yolo = MyYOLO()
        self.model_last_modified = None
        self._load_model()
        # 初始化Flask应用
        self.app = Flask(__name__)
        self._setup_routes()

    def _setup_routes(self):
        """设置Flask路由"""
        self.app.add_url_rule('/detect', 'detect', self._flask_detect, methods=['POST'])
        self.app.add_url_rule('/', 'index', self._index, methods=['GET'])

    def _index(self):
        """首页路由"""
        return "YOLO Detector Service is running"

    def _flask_detect(self):
        """Flask检测接口"""
        try:
            # 检查是否有上传文件
            if 'image' not in request.files:
                return jsonify({'error': 'No image file provided'}), 400
            
            file = request.files['image']
            if file.filename == '':
                return jsonify({'error': 'No image selected'}), 400

            # 读取图片
            image_stream = file.read()
            image_array = np.frombuffer(image_stream, np.uint8)
            image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
            
            if image is None:
                return jsonify({'error': 'Invalid image format'}), 400

            # 执行检测
            annotated_image, results = self.yolo.detect_image_from_array(image)
            
            # 提取边界框信息
            boxes = []
            if hasattr(results[0], 'boxes') and results[0].boxes is not None:
                boxes_data = results[0].boxes
                for box in boxes_data:
                    x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                    conf = box.conf.cpu().numpy()[0]
                    cls = int(box.cls.cpu().numpy()[0])
                    class_name = self.yolo.class_names[cls] if self.yolo.class_names else f"Class {cls}"
                    
                    boxes.append({
                        'x1': int(x1),
                        'y1': int(y1),
                        'x2': int(x2),
                        'y2': int(y2),
                        'confidence': float(conf),
                        'class_id': cls,
                        'class_name': class_name
                    })

            # 返回边界框数据
            return jsonify({
                'boxes': boxes,
                'count': len(boxes)
            })

        except Exception as e:
            return jsonify({'error': str(e)}), 500

    def _load_model(self):
        """加载模型并记录修改时间"""
        if os.path.exists(self.model_path):
            self.yolo.load_model(self.model_path)
            self.model_last_modified = os.path.getmtime(self.model_path)
        else:
            raise FileNotFoundError(f"模型文件 {self.model_path} 不存在")

    def _recheck_images(self, input_folder, output_folder):
        """
        对文件夹中的图片进行重新检测并保存标签文件
        :param input_folder: 输入图片文件夹路径
        :param output_folder: 输出结果文件夹路径
        """
        print(f"重新检测图片{input_folder},保存标签文件到{output_folder}")
        # 检查输入文件夹
        if not os.path.exists(input_folder):
            print(f"输入文件夹 {input_folder} 不存在")
            return

        # 创建输出文件夹
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)

        # 创建labelimg文件夹用于保存原图和标签文件
        labelimg_folder = os.path.join(output_folder, "labelimg")
        if not os.path.exists(labelimg_folder):
            os.makedirs(labelimg_folder)

        # 支持的图片格式
        supported_formats = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif')

        # 获取文件夹中的所有图片文件
        image_files = [
            f for f in os.listdir(input_folder)
            if f.lower().endswith(supported_formats)
        ]

        if not image_files:
            print(f"文件夹 {input_folder} 中没有找到支持的图片文件")
            return

        print(f"找到 {len(image_files)} 个图片文件，开始处理...")

        # 处理每个图片文件
        for image_file in image_files:
            try:
                image_path = os.path.join(input_folder, image_file)
                print(f"正在处理: {image_file}")

                # 执行检测
                annotated_image, results = self.yolo.detect_image(image_path)

                # 保存结果图像
                output_path = os.path.join(output_folder,
                                           f"result_{image_file}")
                self.yolo.save_result(annotated_image, output_path)

                # 复制原图到labelimg文件夹
                target_image_path = os.path.join(labelimg_folder, image_file)
                # if os.path.exists(target_image_path):
                #     os.remove(target_image_path)

                # 保存YOLO格式的标签文件
                label_path = os.path.join(
                    labelimg_folder, f"{os.path.splitext(image_file)[0]}.txt")
                # 读取原图获取尺寸信息
                import cv2
                image = cv2.imread(image_path)
                if image is not None:
                    self.yolo.save_yolo_labels(results, label_path,
                                               image.shape[1], image.shape[0])
                    print(f"标签已保存到: {label_path}")
                # 复制而不是移动原图
                if not os.path.exists(target_image_path):
                    import shutil
                    shutil.copy(image_path, target_image_path)
                print(f"结果已保存到: {output_path}")

            except Exception as e:
                print(f"处理图片 {image_file} 时出错: {str(e)}")
                continue
        print(f"处理完成，共处理 {len(image_files)} 个图片文件")

    def _check_model_update(self):
        """检查模型文件是否被修改"""
        if not os.path.exists(self.model_path):
            return False

        current_modified = os.path.getmtime(self.model_path)
        if self.model_last_modified is None or current_modified != self.model_last_modified:
            # 模型文件被修改，需要重新加载
            self.model_last_modified = current_modified
            try:
                self.yolo.load_model(self.model_path)
                print(f"模型文件 {self.model_path} 已更新，重新加载成功")
                self._recheck_images(r"D:\results\labelimg", r"D:\results")
                return True
            except Exception as e:
                print(f"重新加载模型时出错: {e}")
                return False
        return False

    def detect_once(self):
        """执行单次检测"""
        try:
            self._check_model_update()
            self.yolo.detect_single_image(self.image_path, self.output_folder)
        except Exception as e:
            print(f"Error during detection: {e}")

    def run_periodically(self):
        """定期执行检测"""
        while True:
            # 获取当前时间
            now = datetime.datetime.now()
            # 计算下一分钟的时间点
            next_minute = (now + datetime.timedelta(minutes=1)).replace(
                second=0, microsecond=0)
            # 计算等待时间
            sleep_time = (next_minute - now).total_seconds()
            # 等待到下一分钟
            time.sleep(sleep_time)
            # 执行检测
            self.detect_once()
    
    def run_flask(self, host='0.0.0.0', port=5005):
        """运行Flask服务"""
        self.app.run(host=host, port=port, debug=False)


if __name__ == "__main__":
    model_path = r"./models/best.pt"  # 根据实际情况修改模型路径
    image_path = r"D:\myGitee\control-net\data\onecloud\temp_2.jpg"  # 根据实际情况修改文件夹路径
    output_folder = r"D:\results"  # 结果保存文件夹

    detector = YOLODetector(model_path, image_path, output_folder)

    # 修改为同时运行Flask服务和定期检测
      
    # 在后台线程中运行定期检测
    periodic_thread = threading.Thread(target=detector.run_periodically)
    periodic_thread.daemon = True  # 设置为守护线程，主线程结束时自动结束
    periodic_thread.start()
    
    # 在主线程中运行Flask服务
    detector.run_flask()
